from langchain_core.messages import HumanMessage
from com.wp.langchain.exercise.LLMInit import LLMInit
from langchain_core.chat_history import (
    BaseChatMessageHistory,
    InMemoryChatMessageHistory,
)
from langchain_core.runnables.history import RunnableWithMessageHistory

llm = LLMInit().get_llm()
store = {}

def get_session_history(session_id: str) -> BaseChatMessageHistory:
    if session_id not in store:
        store[session_id] = InMemoryChatMessageHistory()
    return store[session_id]

with_message_history = RunnableWithMessageHistory(llm, get_session_history)

config = {"configurable": {"session_id": "abc2"}}
history_invoke = with_message_history.invoke([HumanMessage(content="Hi! I'm WuPeng")], config=config,)
print(history_invoke.content)

response = with_message_history.invoke([HumanMessage(content="What's my name?")], config=config,)

print(response.content)
